
"""
网络文本主题分类: WebOrganizer/TopicClassifier
共24个类别
"""
import os
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"

from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer
)
from threading import RLock


class TextClassifier(object):

    lock = RLock()
    def __new__(cls, *args, **kwargs):
        with TextClassifier.lock:
            if not hasattr(TextClassifier, "_instance"):
                TextClassifier._instance = object.__new__(cls)
                repo_id = "WebOrganizer/TopicClassifier"
                cls.tokenizer = AutoTokenizer.from_pretrained(repo_id)
                model = AutoModelForSequenceClassification.from_pretrained(repo_id, trust_remote_code=True, use_memory_efficient_attention=False)
                cls.model = model
                cls.class_mapping = model.config.id2label
        return TextClassifier._instance

    def predict(self, text):
        tokens = self.tokenizer(text, return_tensors='pt')
        outputs = self.model(**tokens)

        probs = outputs.logits.softmax(dim=-1)
        idx = probs.argmax(dim=-1).numpy().tolist()[0]
        return self.class_mapping[idx]


if __name__ == '__main__':
    classifier = TextClassifier()
    text = "It is great to see athletes promoting awareness for climate change."
    class_name = classifier.predict(text)
    print(class_name)